from typing import Any, Dict, List, Optional
import json
import time
import redis
from ..base import BaseMemoryModule


class RedisMemoryStore:
    """Redis记忆存储"""
    def __init__(self, host: str, port: int, db: int = 0):
        self.redis = redis.Redis(
            host=host,
            port=port,
            db=db,
            decode_responses=True
        )

    async def store(self, key: str, data: Dict[str, Any], ttl: Optional[int] = None) -> None:
        """存储数据"""
        serialized = json.dumps(data)
        await self.redis.set(key, serialized)
        if ttl:
            await self.redis.expire(key, ttl)

    async def retrieve(self, key: str) -> Optional[Dict[str, Any]]:
        """检索数据"""
        data = await self.redis.get(key)
        return json.loads(data) if data else None

    async def delete(self, key: str) -> None:
        """删除数据"""
        await self.redis.delete(key)

    async def list_keys(self, pattern: str) -> List[str]:
        """列出匹配的键"""
        return await self.redis.keys(pattern)


class DialogueHistory:
    """对话历史管理"""
    def __init__(self, window_size: int = 10):
        self.window_size = window_size
        self.history: List[Dict[str, Any]] = []

    def add(self, entry: Dict[str, Any]) -> None:
        """添加对话记录"""
        self.history.append(entry)
        if len(self.history) > self.window_size:
            self.history.pop(0)

    def get_recent(self, k: int = None) -> List[Dict[str, Any]]:
        """获取最近的对话记录"""
        k = k or self.window_size
        return self.history[-k:]

    def clear(self) -> None:
        """清空历史"""
        self.history.clear()


class MemoryModule(BaseMemoryModule):
    """记忆模块实现"""
    def __init__(self):
        self.store = None
        self.dialogue_history = None
        self.config = None

    def initialize(self, config: Dict[str, Any]) -> None:
        """初始化模块"""
        self.config = config
        self.store = RedisMemoryStore(
            host=config["host"],
            port=config["port"]
        )
        self.dialogue_history = DialogueHistory(
            window_size=config["window_size"]
        )

    def validate_config(self, config: Dict[str, Any]) -> bool:
        """验证配置"""
        required = ["host", "port", "window_size", "ttl", "max_history"]
        return all(key in config for key in required)

    async def store(self, session_id: str, data: Dict[str, Any]) -> None:
        """存储对话历史"""
        if not self.store:
            raise RuntimeError("Module not initialized")

        # 添加时间戳
        data["timestamp"] = time.time()
        
        # 存储到Redis
        key = f"dialogue:{session_id}:{data['timestamp']}"
        await self.store.store(
            key=key,
            data=data,
            ttl=self.config["ttl"]
        )
        
        # 更新对话窗口
        self.dialogue_history.add(data)

    async def retrieve(self, session_id: str, k: int = 5) -> List[Dict[str, Any]]:
        """检索对话历史"""
        if not self.store:
            raise RuntimeError("Module not initialized")

        # 获取所有相关的键
        pattern = f"dialogue:{session_id}:*"
        keys = await self.store.list_keys(pattern)
        
        # 按时间戳排序
        history = []
        for key in sorted(keys, reverse=True)[:k]:
            if data := await self.store.retrieve(key):
                history.append(data)
        
        return history

    async def update(self, session_id: str, data: Dict[str, Any]) -> None:
        """更新对话历史"""
        if not self.store:
            raise RuntimeError("Module not initialized")

        # 获取最新的对话记录
        pattern = f"dialogue:{session_id}:*"
        keys = await self.store.list_keys(pattern)
        
        if keys:
            # 更新最新的记录
            latest_key = max(keys)
            latest_data = await self.store.retrieve(latest_key)
            if latest_data:
                latest_data.update(data)
                await self.store.store(
                    key=latest_key,
                    data=latest_data,
                    ttl=self.config["ttl"]
                )
                
                # 更新对话窗口
                self.dialogue_history.history[-1].update(data)

    async def get_context_window(self) -> List[Dict[str, Any]]:
        """获取当前上下文窗口"""
        return self.dialogue_history.get_recent()

    async def clear_history(self, session_id: str) -> None:
        """清空对话历史"""
        if not self.store:
            raise RuntimeError("Module not initialized")

        # 删除Redis中的记录
        pattern = f"dialogue:{session_id}:*"
        keys = await self.store.list_keys(pattern)
        for key in keys:
            await self.store.delete(key)
        
        # 清空对话窗口
        self.dialogue_history.clear()

    async def prune_old_records(self, session_id: str) -> None:
        """清理旧记录"""
        if not self.store:
            raise RuntimeError("Module not initialized")

        pattern = f"dialogue:{session_id}:*"
        keys = await self.store.list_keys(pattern)
        
        # 保留最新的max_history条记录
        if len(keys) > self.config["max_history"]:
            sorted_keys = sorted(keys)
            old_keys = sorted_keys[:-self.config["max_history"]]
            for key in old_keys:
                await self.store.delete(key) 